from Simulation_with_TDD import *
from TDD.TDD_show import TDD_show
from qiskit.tools.visualization import circuit_drawer
from TDD.TDD import get_node_set, Node, get_int_key

travel = []

def eq_of_nodes(self, other):
    """decide whether the two nodes have the same information"""
    if self.key == other.key and self.idx == other.idx and self.out_weight == other.out_weight:
        return True
    else:
        return False

def get_node_list(node,node_list=list()):
    """Only been used when counting the node number of a TDD"""
    node_list.append(node)
    for k in range(2):
        if node.successor[k]:
            node_list = get_node_list(node.successor[k],node_list)
    return node_list

def pattern_check(node):
    """to determine which pattern the node v belongs to"""
    '''pattern:
    0: not 1/2/3/4
    1: w0(v)=0 and w0(high(v))=0
    2: w1(v)=0 and w1(low(v))=0
    3: w0(v)=w1(v)=1 and w1(low(v))=0 and w0(high(v))=0 and high(low(v))=low(high(v))
    4: w0(v)=w1(v)=1 and w1(low(v))=0 and w0(high(v))=0 and high(low(v))≠low(high(v))
    '''
    if node.key == 1 or node.key[0] == 'y':
        return 0
    if get_int_key(node.out_weight[0]) == (0, 0):
        if get_int_key(node.successor[1].out_weight[0]) != (0, 0):
            return 0
        else:
            return 1
    if get_int_key(node.out_weight[1]) == (0, 0):
        if get_int_key(node.successor[0].out_weight[1]) != (0, 0):
            return 0
        else:
            return 2
    if node.successor[0].key == node.successor[1].key:
        if get_int_key(node.successor[0].out_weight[1]) != (0, 0) or \
                get_int_key(node.successor[1].out_weight[0]) != (0, 0):
            return 0
        else:
            if get_int_key(node.out_weight[0]) == get_int_key(node.out_weight[1]) and \
                    node.successor[0].successor[0] == node.successor[1].successor[1]:
                return 3
            else:
                return 4
    return 0



def RR5(tdd, cur):
    """to return a tdd that removes redundant nodes and preserves circuit equivalence"""
    '''preorder traversal on the graph'''
    global travel
    if cur.key == 1:  # 叶子节点
        return cur
    pattern = pattern_check(cur)
    if pattern == 1:
        cur.key = cur.successor[1].key
        cur.successor[1] = cur.successor[1].successor[1]
        cur.out_weight[1] = 1
    if pattern == 2:
        cur.key = cur.successor[0].key
        cur.successor[0] = cur.successor[0].successor[0]
        cur.out_weight[0] = 1
    if pattern == 3:
        cur =  cur.successor[0].successor[0]
        cur = RR5(tdd,cur)
        return cur
    if pattern == 4:
        cur.key = cur.successor[0].key
        cur.successor[0] = cur.successor[0].successor[0]
        cur.successor[1] = cur.successor[1].successor[1]

    for k in range(2):
        cur.successor[k] = RR5(tdd,cur.successor[k])
    return cur



def equivalence_check(tdd1,tdd2):
    """decide whether the two tdds are isomorphic"""
    node_list1 = list()
    node_list2 = list()
    node_list1 = get_node_list(tdd1.node,node_list1)
    node_list2 = get_node_list(tdd2.node,node_list2)
    len1 = len(node_list1)
    len2 = len(node_list2)
    print('Node number of tdd1 is ', len1)
    print('Node number of tdd2 is ', len2)
    if len1 != len2:
        return False
    for i in range(len1):
        if not eq_of_nodes(node_list1[i],node_list2[i]):
            return False
    return True
    # if tdd1 == tdd2:
    #     return True
    # else:
    #     return False

#global variables
node_set = set()


# def deCont(tdd):
#     global node_set
#     #收集节点信息
#     node_set = get_node_set(tdd.node)
#     newNode_list = list()             #记录新加的节点
#     #遍历set录入入度信息
#     for cur in node_set:
#         #print("节点位置信息对比：",cur.key," at ",cur)
#         #对每个节点： 如果他的 child 的 in_weight == -1 则表示还没有赋值 child 的in weight
#         for k in range(2):
#             if cur.successor[k] :
#                 child = cur.successor[k]
#                 if not child.key == 1:
#                     if child.in_weight == -99999:
#                         #print("孩子节点地址信息：",child.key," at ",child)
#                                             #还未赋值
#                         #print("对节点",child.key,"正在初次赋值in weight")
#                         child.in_weight = cur.out_weight[k]
#                     elif child.in_weight == cur.out_weight[k]:       #已赋值，入度一致
#                         #print("当前节点：",child.key,"child的inweight：",child.in_weight, "父节点outweight：",cur.out_weight[k])
#                         continue
#                     else:                                           #已赋值，且入度不一致，检查并构造新节点
#                         #print("对节点",child.key,"进行分裂")
#                         newNode = Node(child.key)
#                         newNode.in_weight = cur.out_weight[k]
#                         for kk in range(2):
#                             newNode.out_weight[kk] = child.out_weight[kk]
#                             newNode.successor[kk] = child.successor[kk]
#                         if newNode in newNode_list:
#                             newNode = newNode_list[newNode_list.index(newNode)]
#                         else:
#                             newNode_list.append(newNode)
#                         cur.successor[k] = newNode


def RR5_and_eqCheck():
    path='eq/'
    file_list = os.listdir(path)
    tdd = [None]*2
    flag = 0
    print("check len:",len(tdd))
    #tensor TDD构造过程
    for file_name in file_list:
        print("check flag:",flag)
        if flag >= 2:
            print(file_name,"文件夹里只能放两个文件作比较")
        print('circuit：',file_name)
        try:
            cir,res = CreateDGfromQASMfile(file_name, path, flag_single=True)
        except:
            continue
        dag_cir=res[0]

        num_qubit = get_real_qubit_num(dag_cir)
        print('qubits:',num_qubit)
        gate_num = get_gates_number(dag_cir)
        print('gates number:',gate_num)

        #try:
        #t_start = time.time()
        tmpTDD,max_node_num,block_num=Simulation_with_cir_partion(dag_cir, num_qubit,0,num_qubit//2,num_qubit//2+1)
        tdd[flag] = tmpTDD
        #run_time=time.time()-t_start
        #print('Run time:',run_time)
        print('Max node num:',max_node_num)
        print('Final node number:',tmpTDD.node_number())
        #except:
        #    print('Time out!')

        tdd[flag].node = RR5(tdd[flag],tdd[flag].node)                  #RR5

        TDD_show(tdd[flag],file_name+"_out")
        print(cir.draw())
        print('----------------')
        flag+=1
    print('=================')
  
    #eqCheck
    print(equivalence_check(tdd[0],tdd[1]))


if __name__ == '__main__':
    
    RR5_and_eqCheck()